/*
 *  Copyright 2010 INFN - APE group
 *
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 */

#pragma once

#include <cuqu/config.h>

namespace cuqu {

    namespace detail {

        struct barrier_tree_based {
            enum { numGroups = 6, maxGroups = 20 };
            atomic_t m_cnt1[maxGroups];   // level 1
            atomic_t m_cnt2;              // level 2
        };

#define CUQU_INIT_BARRIER_TREE_BASED { { 0 }, 0 }

        static __inline__ __device__ int barrier_wait(barrier_tree_based *gbl_barrier, int n_waiters) {

            // Thread ID and block ID.
            const unsigned int tid = threadIdx.x + blockDim.x * threadIdx.y + blockDim.x * blockDim.y * threadIdx.z;
            const unsigned int bid = blockIdx.x + gridDim.x * blockIdx.y;

            // Total numbers of groups.
            //unsigned int numGroups = barrier_tree_based::numGroups;
            unsigned int numGroups = (unsigned int) ceilf(sqrtf(n_waiters));
            
            // The group index of the block.
            unsigned int groupIdx;

            // Check if the number of waiters is less than number of groups.
            if (numGroups > n_waiters || numGroups > barrier_tree_based::maxGroups) {
                numGroups = n_waiters;
            }
             
            int cnt;
            int retry = 200000;

            // Default return code value.
            //int retcode = SUCCESS;

            // Only thread zero of each block compute the barrier.
            if (tid == 0) {
                
                // -----------------------
                // First level of the tree
                // -----------------------
                
                groupIdx = bid % numGroups;

                // Calculare the number of blocks in the group.
                unsigned int numBlocksInGroup = n_waiters / numGroups;
               
                // Does numGroups divide n_waiters? 
                if (groupIdx < (n_waiters % numGroups)) {
                    // If false then some group have one block more.
                    numBlocksInGroup++;
                }
    
                // Incremente the group atomic counter.
                unsigned int oldcnt = atomic_mod_inc(&gbl_barrier->m_cnt1[groupIdx], numBlocksInGroup - 1);

#if defined(CUQU_DEBUG)
                if(oldcnt >= numBlocksInGroup) {
                    cuqu_error("ERROR invalid barrier counter %d\n", oldcnt);
                }
#endif
                // Begin the infernal loop.
                while (true) {

                    // Read the atomic counter of the group groupIdx.
                    cnt = atomic_read(&gbl_barrier->m_cnt1[groupIdx]);
                    cuqu_trace("cnt: %d\n", cnt);

                    // Have all blocks incremented the counter?  
                    if(cnt == 0) {
                        cuqu_trace("[%d:%d:%d] first level barrier %p done, counter reached %d\n", groupIdx, bid, tid, gbl_barrier, cnt);
                        break; // Continue to second level.
                    }

#if defined(CUQU_DEBUG)
                    // Has the counter passed the number of blocks in group?
                    else if(cnt >= numBlocksInGroup) {
                        cuqu_error("[%d:%d] ERROR invalid barrier counter %d\n", bid, tid, cnt);
                        __syncthreads();
                        return INVALID;
                        //retcode = INVALID; // Exit from the infernal loop with error.
                        //break;
                    }
#endif
                    --retry;
                    if(retry < 0) {
                        cuqu_error("[%d:%d:%d] ERROR barrier %p timeout polling counter, n_waiters=%d cnt=%d numBlocksInGroup=%d\n", groupIdx, bid, tid, gbl_barrier, n_waiters, cnt, numBlocksInGroup);
                        __syncthreads();
                        return WOULDBLOCK;
                        //retcode = WOULDBLOCK;
                        //break;
                    }
                }
            }

            // All threads in block are synchronized.
            __syncthreads();


            // Reset some variables.
            retry = 200000;


            // Only thread zero of each block compute the barrier.
            if (tid == 0) {
                
                // ------------------------
                // Second level of the tree
                // ------------------------

                // Only the first thread of the first numGroups groups increment second level atomic couter variable.
                if (bid < numGroups) {
                    // Increment second level atomic variable.
                    unsigned int oldcnt = atomic_mod_inc(&gbl_barrier->m_cnt2, numGroups - 1);
#if defined(CUQU_DEBUG)
                    if(oldcnt >= numGroups) {
                        cuqu_error("ERROR invalid barrier counter %d\n", oldcnt);
                    }
#endif
                }
                
                // Infernal loop.
                while (true) {
                    // Read the second level atomic counter.
                    cnt = atomic_read(&gbl_barrier->m_cnt2);

                    // Have all groups incremented the counter?  
                    if(cnt == 0) {
                        cuqu_trace("[%d:%d:%d] second level barrier %p done, counter reached %d\n", 
                                    groupIdx, bid, tid, gbl_barrier, cnt);
                        break;
                    }

#if defined(CUQU_DEBUG)
                    // Has the counter passed the number of blocks in group?
                    else if(cnt >= numGroups) {
                        cuqu_error("[%d:%d] ERROR invalid barrier counter %d\n", bid, tid, cnt);
                        __syncthreads();
                        return INVALID;
                        //retcode = INVALID; // Exit from the infernal loop with error.
                        //break;
                    }
#endif

                    --retry;
                    if(retry < 0) {
                        cuqu_error("[%d:%d:%d] ERROR second barrier %p timeout polling counter, cnt=%d numGroups=%d\n",
                                    groupIdx, bid, tid, gbl_barrier, cnt, numGroups);
                        __syncthreads();
                        return WOULDBLOCK;
                        //retcode = WOULDBLOCK;
                        //break;
                    }
                }
            }

            // All threads in the blocks are synchronized.
            __syncthreads();
            return SUCCESS;
        }

    } // end namespace detail
    
} // end namespace cuqu
